#!/usr/bin/env python3
"""
EchoKit Message Handler - 处理EchoKit服务器消息格式转换
解决LLM返回格式与机器人客户端期望格式不匹配的问题

支持MessagePack和JSON两种格式的消息解析。
"""

import re
import json
import logging
import time
from typing import Dict, Any, Optional, Tuple, Union
from geometry_msgs.msg import Twist

try:
    import msgpack
    MSGPACK_AVAILABLE = True
except ImportError:
    MSGPACK_AVAILABLE = False


class EchoKitMessageHandler:
    """处理EchoKit服务器消息的格式转换和命令解析"""
    
    def __init__(self, logger=None):
        self.logger = logger or logging.getLogger(__name__)
        
        # 命令映射表 - 将方括号命令映射到具体动作
        self.command_mapping = {
            'forward': {'linear': 0.3, 'angular': 0.0},
            'backward': {'linear': -0.3, 'angular': 0.0},
            'turn_left': {'linear': 0.0, 'angular': 1.0},
            'turn_right': {'linear': 0.0, 'angular': -1.0},
            'left': {'linear': 0.0, 'angular': 1.0},
            'right': {'linear': 0.0, 'angular': -1.0},
            'stop': {'linear': 0.0, 'angular': 0.0},
        }
        
        # 视觉查询关键词映射
        self.vision_query_keywords = {
            '看到': 'what_do_you_see',
            '看见': 'what_do_you_see', 
            '观察': 'observe',
            '识别': 'identify',
            '分析': 'analyze',
            '描述': 'describe',
            '这是什么': 'what_is_this',
            '那是什么': 'what_is_that',
            'what do you see': 'what_do_you_see',
            'what is this': 'what_is_this',
            'describe': 'describe',
            'analyze': 'analyze',
            'identify': 'identify'
        }
        
        # 方括号命令的正则表达式
        self.bracket_pattern = re.compile(r'\[([^\]]+)\]')
        
    def parse_message(self, raw_message: Union[str, bytes]) -> Optional[Dict[str, Any]]:
        """
        解析来自EchoKit服务器的原始消息
        
        Args:
            raw_message: 原始消息（字符串或字节）
            
        Returns:
            解析后的消息字典，如果解析失败返回None
        """
        try:
            # 如果是字节数据，尝试MessagePack解析
            if isinstance(raw_message, bytes):
                return self._handle_messagepack_message(raw_message)
            
            # 如果是字符串，尝试JSON解析
            if isinstance(raw_message, str):
                # 首先尝试解析为JSON
                try:
                    json_data = json.loads(raw_message)
                    return self._handle_json_message(json_data)
                except json.JSONDecodeError:
                    pass
                
                # 如果不是JSON，尝试解析方括号格式
                bracket_match = self.bracket_pattern.search(raw_message)
                if bracket_match:
                    command = bracket_match.group(1).lower().strip()
                    return self._handle_bracket_command(command, raw_message)
                
                # 如果都不匹配，尝试直接文本解析
                return self._handle_text_message(raw_message)
            
            if self.logger:
                self.logger.warn(f"Unsupported message type: {type(raw_message)}")
            return None
            
        except Exception as e:
            if self.logger:
                self.logger.error(f"Error parsing message: {e}")
            return None
    
    def _handle_messagepack_message(self, msgpack_data: bytes) -> Optional[Dict[str, Any]]:
        """处理MessagePack格式的消息"""
        if not MSGPACK_AVAILABLE:
            if self.logger:
                self.logger.error("MessagePack not available. Install with: pip install msgpack")
            return None
        
        try:
            # 解析MessagePack数据
            data = msgpack.unpackb(msgpack_data, raw=False)
            
            if self.logger:
                self.logger.info(f"Parsed MessagePack data: {data}")
            
            # 处理EchoKit ServerEvent格式
            if isinstance(data, dict):
                return self._handle_server_event(data)
            elif isinstance(data, list) and len(data) >= 2:
                # MessagePack可能使用数组格式 [event_type, data]
                event_type = data[0]
                event_data = data[1] if len(data) > 1 else {}
                return self._handle_server_event({event_type: event_data})
            else:
                if self.logger:
                    self.logger.warn(f"Unknown MessagePack format: {data}")
                return None
                
        except Exception as e:
            if self.logger:
                self.logger.error(f"Error parsing MessagePack: {e}")
            return None
    
    def _handle_server_event(self, event_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理EchoKit ServerEvent格式的消息"""
        # 处理Vision事件 (视觉交互消息)
        if 'VisionQuery' in event_data:
            return self._handle_vision_query_event(event_data['VisionQuery'])
        
        if 'VisionResponse' in event_data:
            return self._handle_vision_response_event(event_data['VisionResponse'])
            
        if 'CameraStatus' in event_data:
            return self._handle_camera_status_event(event_data['CameraStatus'])
        
        # 处理Action事件 (最重要的机器人命令)
        if 'Action' in event_data:
            action_data = event_data['Action']
            
            # 处理数组格式: {'Action': ['forward']}
            if isinstance(action_data, list) and len(action_data) > 0:
                action_str = str(action_data[0]).lower().strip()
                return {
                    'type': 'voice_command',
                    'command': action_str,
                    'action': self.command_mapping.get(action_str),
                    'original_message': event_data,
                    'format': 'messagepack_action_array'
                }
            
            # 处理字典格式: {'Action': {'action': 'forward'}}
            elif isinstance(action_data, dict) and 'action' in action_data:
                action_str = action_data['action'].lower().strip()
                return {
                    'type': 'voice_command',
                    'command': action_str,
                    'action': self.command_mapping.get(action_str),
                    'original_message': event_data,
                    'format': 'messagepack_action_dict'
                }
            
            # 处理字符串格式: {'Action': 'forward'} 或 {'Action': '{"type": "timed_motion", ...}'}
            elif isinstance(action_data, str):
                action_str = action_data.strip()
                
                # 检查是否是JSON格式的结构化命令
                if action_str.startswith('{') and action_str.endswith('}'):
                    try:
                        json_command = json.loads(action_str)
                        if self.logger:
                            self.logger.info(f"Parsed JSON action command: {json_command}")
                        
                        # 直接返回结构化命令
                        return {
                            'type': 'structured_command',
                            'command_data': json_command,
                            'original_message': event_data,
                            'format': 'messagepack_json_action'
                        }
                    except json.JSONDecodeError as e:
                        if self.logger:
                            self.logger.warn(f"Failed to parse JSON action: {e}")
                
                # 处理简单字符串命令
                action_lower = action_str.lower().strip()
                return {
                    'type': 'voice_command',
                    'command': action_lower,
                    'action': self.command_mapping.get(action_lower),
                    'original_message': event_data,
                    'format': 'messagepack_action_string'
                }
        
        # 处理ASR事件 (语音识别结果)
        if 'ASR' in event_data:
            asr_data = event_data['ASR']
            if isinstance(asr_data, dict) and 'text' in asr_data:
                text = asr_data['text']
                # 尝试从ASR文本中提取命令
                bracket_match = self.bracket_pattern.search(text)
                if bracket_match:
                    command = bracket_match.group(1).lower().strip()
                    return {
                        'type': 'voice_command',
                        'command': command,
                        'action': self.command_mapping.get(command),
                        'original_message': event_data,
                        'format': 'messagepack_asr'
                    }
                else:
                    return {
                        'type': 'asr_result',
                        'text': text,
                        'original_message': event_data,
                        'format': 'messagepack_asr'
                    }
        
        # 处理其他事件类型
        event_types = ['StartAudio', 'EndAudio', 'AudioChunk', 'EndResponse', 'HelloStart', 'HelloEnd']
        for event_type in event_types:
            if event_type in event_data:
                return {
                    'type': event_type.lower(),
                    'data': event_data[event_type],
                    'original_message': event_data,
                    'format': 'messagepack_event'
                }
        
        # 未知事件类型
        return {
            'type': 'unknown',
            'original_message': event_data,
            'format': 'messagepack_unknown'
        }
    
    def _handle_vision_query_event(self, vision_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理视觉查询事件"""
        query_text = vision_data.get('query', '')
        image_data = vision_data.get('image_data')
        query_id = vision_data.get('query_id', str(int(time.time() * 1000)))
        
        return {
            'type': 'vision_query',
            'query_id': query_id,
            'query_text': query_text,
            'image_data': image_data,
            'timestamp': vision_data.get('timestamp', int(time.time() * 1000)),
            'original_message': vision_data,
            'format': 'vision_query'
        }
    
    def _handle_vision_response_event(self, response_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理视觉响应事件"""
        return {
            'type': 'vision_response',
            'query_id': response_data.get('query_id', ''),
            'response_text': response_data.get('response', ''),
            'confidence': response_data.get('confidence', 0.0),
            'processing_time': response_data.get('processing_time', 0),
            'model_used': response_data.get('model_used', 'unknown'),
            'timestamp': response_data.get('timestamp', int(time.time() * 1000)),
            'original_message': response_data,
            'format': 'vision_response'
        }
    
    def _handle_camera_status_event(self, status_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理摄像头状态事件"""
        return {
            'type': 'camera_status',
            'connected': status_data.get('connected', False),
            'resolution': status_data.get('resolution', ''),
            'fps': status_data.get('fps', 0),
            'format': status_data.get('format', ''),
            'stream_quality': status_data.get('stream_quality', 0.0),
            'timestamp': status_data.get('timestamp', int(time.time() * 1000)),
            'original_message': status_data,
            'format': 'camera_status'
        }

    def _handle_json_message(self, json_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理JSON格式的消息"""
        # 如果JSON中包含命令信息
        if 'command' in json_data:
            command = json_data['command'].lower().strip()
            return {
                'type': 'voice_command',
                'command': command,
                'action': self.command_mapping.get(command),
                'original_message': json_data,
                'format': 'json'
            }
        
        # 如果JSON中包含文本内容，尝试从中提取命令
        if 'content' in json_data or 'text' in json_data:
            text = json_data.get('content', json_data.get('text', ''))
            bracket_match = self.bracket_pattern.search(text)
            if bracket_match:
                command = bracket_match.group(1).lower().strip()
                return {
                    'type': 'voice_command',
                    'command': command,
                    'action': self.command_mapping.get(command),
                    'original_message': json_data,
                    'format': 'json_with_bracket'
                }
        
        # 返回原始JSON消息
        return {
            'type': json_data.get('type', 'unknown'),
            'original_message': json_data,
            'format': 'json'
        }
    
    def _handle_bracket_command(self, command: str, raw_message: str) -> Dict[str, Any]:
        """处理方括号格式的命令"""
        return {
            'type': 'voice_command',
            'command': command,
            'action': self.command_mapping.get(command),
            'original_message': raw_message,
            'format': 'bracket'
        }
    
    def _handle_text_message(self, text: str) -> Dict[str, Any]:
        """处理纯文本消息，尝试从中识别命令"""
        text_lower = text.lower().strip()
        
        # 检查是否是视觉查询
        for keyword, query_type in self.vision_query_keywords.items():
            if keyword in text_lower:
                return {
                    'type': 'vision_query',
                    'query_id': str(int(time.time() * 1000)),
                    'query_text': text,
                    'query_type': query_type,
                    'timestamp': int(time.time() * 1000),
                    'original_message': text,
                    'format': 'text_vision_query'
                }
        
        # 尝试匹配已知运动命令
        for cmd, action in self.command_mapping.items():
            if cmd in text_lower:
                return {
                    'type': 'voice_command',
                    'command': cmd,
                    'action': action,
                    'original_message': text,
                    'format': 'text'
                }
        
        # 如果没有匹配到命令，返回文本消息
        return {
            'type': 'text_message',
            'content': text,
            'original_message': text,
            'format': 'text'
        }
    
    def create_twist_message(self, action: Dict[str, float]) -> Twist:
        """根据动作字典创建Twist消息"""
        twist = Twist()
        if action:
            twist.linear.x = action.get('linear', 0.0)
            twist.angular.z = action.get('angular', 0.0)
        return twist
    
    def is_valid_command(self, command: str) -> bool:
        """检查命令是否有效"""
        return command.lower().strip() in self.command_mapping
    
    def get_available_commands(self) -> list:
        """获取所有可用命令列表"""
        return list(self.command_mapping.keys())
    
    def parse_voice_command(self, message: Union[str, bytes]) -> Optional[Dict[str, Any]]:
        """
        解析语音命令的便捷方法
        
        这是为了兼容测试代码和其他调用者而提供的便捷接口。
        实际上调用的是 parse_message 方法。
        
        Args:
            message: 输入消息（字符串或字节）
            
        Returns:
            解析后的命令字典，包含action字段，如果解析失败返回None
        """
        try:
            result = self.parse_message(message)
            
            if result is None:
                return None
            
            # 如果解析结果包含action信息，直接返回
            if result.get('type') == 'voice_command' and 'action' in result:
                return result
            
            # 如果是文本消息，尝试从中提取命令
            if result.get('type') == 'text_message':
                content = result.get('content', '').lower().strip()
                
                # 中文命令映射
                chinese_commands = {
                    '前进': 'forward',
                    '后退': 'backward', 
                    '左转': 'turn_left',
                    '右转': 'turn_right',
                    '停止': 'stop',
                    '向前': 'forward',
                    '向后': 'backward',
                    '左': 'turn_left',
                    '右': 'turn_right',
                    '停': 'stop'
                }
                
                # 检查中文命令
                for chinese, english in chinese_commands.items():
                    if chinese in content:
                        return {
                            'type': 'voice_command',
                            'command': english,
                            'action': english,
                            'original_message': message,
                            'format': 'chinese_text'
                        }
                
                # 检查英文命令
                for cmd in self.command_mapping.keys():
                    if cmd in content:
                        return {
                            'type': 'voice_command',
                            'command': cmd,
                            'action': cmd,
                            'original_message': message,
                            'format': 'english_text'
                        }
            
            # 如果是结构化命令，尝试提取action
            if result.get('type') == 'structured_command':
                command_data = result.get('command_data', {})
                if 'action' in command_data:
                    action = command_data['action']
                    return {
                        'type': 'voice_command',
                        'command': action,
                        'action': action,
                        'original_message': message,
                        'format': 'structured'
                    }
            
            # 返回原始结果
            return result
            
        except Exception as e:
            if self.logger:
                self.logger.error(f"Error in parse_voice_command: {e}")
            return None
    
    def format_response_for_echokit(self, message_type: str, content: str) -> str:
        """格式化发送给EchoKit服务器的响应消息"""
        response = {
            'type': message_type,
            'content': content,
            'timestamp': int(time.time() * 1000)
        }
        return json.dumps(response)
    
    def create_vision_query_message(self, query_text: str, image_data: Optional[bytes] = None) -> Dict[str, Any]:
        """创建视觉查询消息"""
        return {
            'VisionQuery': {
                'query_id': str(int(time.time() * 1000)),
                'query': query_text,
                'image_data': image_data,
                'timestamp': int(time.time() * 1000)
            }
        }
    
    def create_vision_response_message(self, query_id: str, response_text: str, 
                                     confidence: float = 1.0, processing_time: int = 0,
                                     model_used: str = 'unknown') -> Dict[str, Any]:
        """创建视觉响应消息"""
        return {
            'VisionResponse': {
                'query_id': query_id,
                'response': response_text,
                'confidence': confidence,
                'processing_time': processing_time,
                'model_used': model_used,
                'timestamp': int(time.time() * 1000)
            }
        }
    
    def create_camera_status_message(self, connected: bool, resolution: str = '', 
                                   fps: int = 0, format_type: str = '',
                                   stream_quality: float = 0.0) -> Dict[str, Any]:
        """创建摄像头状态消息"""
        return {
            'CameraStatus': {
                'connected': connected,
                'resolution': resolution,
                'fps': fps,
                'format': format_type,
                'stream_quality': stream_quality,
                'timestamp': int(time.time() * 1000)
            }
        }
    
    def is_vision_query(self, text: str) -> bool:
        """检查文本是否是视觉查询"""
        text_lower = text.lower().strip()
        return any(keyword in text_lower for keyword in self.vision_query_keywords.keys())


# 测试函数
def test_message_handler():
    """测试消息处理器的功能"""
    handler = EchoKitMessageHandler()
    
    # 测试用例 - 字符串消息
    test_messages = [
        '[forward]',
        '[turn_left]',
        '{"type": "command", "command": "forward"}',
        '{"content": "[stop]"}',
        'forward',
        'invalid message',
        '[unknown_command]'
    ]
    
    print("Testing EchoKit Message Handler - parse_message:")
    for msg in test_messages:
        result = handler.parse_message(msg)
        print(f"Input: {msg}")
        print(f"Output: {result}")
        print("-" * 50)
    
    # 测试 parse_voice_command 方法
    voice_test_messages = [
        '前进',
        '后退', 
        '左转',
        '右转',
        '停止',
        '[forward]',
        'forward',
        '{"action": "stop"}',
        'unknown command'
    ]
    
    print("\nTesting parse_voice_command method:")
    for msg in voice_test_messages:
        result = handler.parse_voice_command(msg)
        print(f"Input: {msg}")
        print(f"Output: {result}")
        print("-" * 50)
    
    # 测试MessagePack消息（如果可用）
    if MSGPACK_AVAILABLE:
        print("\nTesting MessagePack Messages:")
        
        # 模拟EchoKit ServerEvent消息
        test_msgpack_data = [
            {'Action': {'action': 'forward'}},
            {'ASR': {'text': 'move forward'}},
            {'ASR': {'text': '[turn_left]'}},
            {'StartAudio': {'text': 'Starting audio'}},
            {'EndResponse': {}},
        ]
        
        for data in test_msgpack_data:
            try:
                msgpack_bytes = msgpack.packb(data)
                result = handler.parse_message(msgpack_bytes)
                print(f"Input (MessagePack): {data}")
                print(f"Output: {result}")
                print("-" * 50)
            except Exception as e:
                print(f"Error testing MessagePack data {data}: {e}")
    else:
        print("\nMessagePack not available for testing. Install with: pip install msgpack")


if __name__ == '__main__':
    test_message_handler()